// Copyright 2017 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <fcntl.h>
#include <stdint.h>
#include <stdlib.h>
#include <sys/stat.h>

#include <fbl/function.h>
#include <fbl/string.h>
#include <fbl/string_buffer.h>
#include <fbl/string_printf.h>
#include <fbl/unique_fd.h>
#include <fs-management/mount.h>
#include <fs-test-utils/fixture.h>
#include <fs-test-utils/perftest.h>
#include <perftest/perftest.h>
#include <unittest/unittest.h>

#include <utility>

namespace fs_bench {
namespace {

using fs_test_utils::Fixture;
using fs_test_utils::FixtureOptions;
using fs_test_utils::PerformanceTestOptions;
using fs_test_utils::TestCaseInfo;
using fs_test_utils::TestInfo;

constexpr int kWriteReadCycles = 3;

fbl::String GetBigFilePath(const Fixture& fixture) {
  fbl::String path = fbl::StringPrintf("%s/bigfile.txt", fixture.fs_path().c_str());
  return path;
}

bool WriteBigFile(ssize_t data_size, perftest::RepeatState* state, Fixture* fixture) {
  BEGIN_HELPER;

  fbl::unique_fd fd(open(GetBigFilePath(*fixture).c_str(), O_CREAT | O_WRONLY));
  ASSERT_TRUE(fd);
  state->DeclareStep("write");
  uint8_t data[data_size];
  uint8_t pattern = static_cast<uint8_t>(rand_r(fixture->mutable_seed()) % (1 << 8));
  memset(data, pattern, data_size);

  while (state->KeepRunning()) {
    ASSERT_EQ(write(fd.get(), data, data_size), data_size);
  }

  END_HELPER;
}

bool ReadBigFile(ssize_t data_size, perftest::RepeatState* state, Fixture* fixture) {
  BEGIN_HELPER;
  fbl::unique_fd fd(open(GetBigFilePath(*fixture).c_str(), O_RDONLY));

  uint8_t pattern = static_cast<uint8_t>(rand_r(fixture->mutable_seed()) % (1 << 8));
  ASSERT_TRUE(fd);
  state->DeclareStep("read");
  uint8_t data[data_size];

  while (state->KeepRunning()) {
    ASSERT_EQ(read(fd.get(), data, data_size), data_size);
    ASSERT_EQ(data[0], pattern);
  }

  END_HELPER;
}

constexpr char kBaseComponent[] = "/aaa";

constexpr size_t kComponentLength = fbl::constexpr_strlen(kBaseComponent);

struct PathComponentGen {
  PathComponentGen() { memcpy(current, kBaseComponent, kComponentLength + 1); }

  // Advances current to the next component, following alphabetical order.
  // E.g: /aaa -> /aab ..../aaz -> /aba
  void Next() {
    for (int i = 3; i > 0; --i) {
      char next = static_cast<char>(static_cast<uint8_t>(current[i]) + 1);
      if (next > 'z') {
        current[i] = 'a';
      } else {
        current[i] = next;
        break;
      }
    }
  }
  // Add extra byte for null termination.
  char current[kComponentLength + 1];
};

bool PathWalkDown(const fbl::String& op_name, const fbl::Function<int(const char*)>& op,
                  perftest::RepeatState* state, Fixture* fixture,
                  fbl::StringBuffer<fs_test_utils::kPathSize>* path) {
  BEGIN_HELPER;
  PathComponentGen component;
  path->Append(fixture->fs_path());

  while (state->KeepRunning()) {
    path->Append(component.current);
    ASSERT_EQ(op(path->c_str()), 0);
    component.Next();
  }
  END_HELPER;
}

bool PathWalkUp(const fbl::String& op_name, const fbl::Function<int(const char*)>& op,
                perftest::RepeatState* state, Fixture* fixture,
                fbl::StringBuffer<fs_test_utils::kPathSize>* path) {
  BEGIN_HELPER;
  while (state->KeepRunning() && *path != fixture->fs_path()) {
    ASSERT_EQ(op(path->c_str()), 0, path->c_str());
    uint32_t new_size = static_cast<uint32_t>(path->length() - kComponentLength);
    path->Resize(new_size);
  }
  END_HELPER;
}

// Wrapper so state can be shared across calls.
class PathWalkOp {
 public:
  PathWalkOp() = default;
  PathWalkOp(const PathWalkOp&) = delete;
  PathWalkOp(PathWalkOp&&) = delete;
  PathWalkOp& operator=(const PathWalkOp&) = delete;
  PathWalkOp& operator=(PathWalkOp&&) = delete;
  ~PathWalkOp() = default;

  // Will add components until |state::KeepGoing| returns false.
  bool Mkdir(perftest::RepeatState* state, Fixture* fixture) {
    path_.Clear();
    return PathWalkDown(
        "mkdir", [](const char* path) { return mkdir(path, 0666); }, state, fixture, &path_);
  }

  // Will stat components until |state::KeepGoing| returns false.
  bool Stat(perftest::RepeatState* state, Fixture* fixture) {
    path_.Clear();
    return PathWalkDown(
        "stat",
        [](const char* path) {
          struct stat buff;
          return stat(path, &buff);
        },
        state, fixture, &path_);
  }

  // Will unlink components until |state::KeepGoing| returns false.
  bool Unlink(perftest::RepeatState* state, Fixture* fixture) {
    return PathWalkUp("unlink", unlink, state, fixture, &path_);
  }

 private:
  fbl::StringBuffer<fs_test_utils::kPathSize> path_;
};

}  // namespace

bool RunBenchmark(int argc, char** argv) {
  FixtureOptions f_opts = FixtureOptions::Default(DISK_FORMAT_MINFS);
  PerformanceTestOptions p_opts;
  const int rw_test_sample_counts[] = {
      1024, 2048, 4096, 8192, 16384,
  };

  if (!fs_test_utils::ParseCommandLineArgs(argc, argv, &f_opts, &p_opts)) {
    return false;
  }

  fbl::Vector<TestCaseInfo> testcases;
  // Just do a single cycle for unittest mode.
  int cycles = (p_opts.is_unittest) ? 1 : kWriteReadCycles;
  // Read Write tests.
  for (int test_sample_count : rw_test_sample_counts) {
    TestCaseInfo testcase;
    testcase.sample_count = test_sample_count;
    testcase.name = fbl::StringPrintf("%s/Bigfile/16Kbytes/%d-Ops",
                                      disk_format_string_[f_opts.fs_type], test_sample_count);
    testcase.teardown = false;
    for (int cycle = 0; cycle < cycles; ++cycle) {
      TestInfo write_test, read_test;
      write_test.name = fbl::StringPrintf("%s/%d-Cycle/Write", testcase.name.c_str(), cycle + 1);
      write_test.test_fn = [](perftest::RepeatState* state, Fixture* fixture) {
        return WriteBigFile(16 * (1 << 10), state, fixture);
      };
      write_test.required_disk_space = test_sample_count * 16 * 1024 * (cycle + 1);
      testcase.tests.push_back(std::move(write_test));

      read_test.name = fbl::StringPrintf("%s/%d-Cycle/Read", testcase.name.c_str(), cycle + 1);
      read_test.test_fn = [](perftest::RepeatState* state, Fixture* fixture) {
        return ReadBigFile(16 * (1 << 10), state, fixture);
      };
      read_test.required_disk_space = test_sample_count * 16 * 1024 * (cycle + 1);
      testcase.tests.push_back(std::move(read_test));
    }
    testcases.push_back(std::move(testcase));
  }

  // Path walk tests.
  const int path_walk_sample_counts[] = {
      125,
      250,
      500,
  };

  PathWalkOp pw_op;
  for (int test_sample_count : path_walk_sample_counts) {
    TestCaseInfo testcase;
    testcase.name = fbl::StringPrintf("%s/PathWalk/%d-Components",
                                      disk_format_string_[f_opts.fs_type], test_sample_count);
    testcase.sample_count = test_sample_count;
    testcase.teardown = false;

    TestInfo mkdir_test;
    mkdir_test.name = fbl::StringPrintf("%s/Mkdir", testcase.name.c_str());
    mkdir_test.test_fn = fbl::BindMember(&pw_op, &PathWalkOp::Mkdir);
    testcase.tests.push_back(std::move(mkdir_test));

    TestInfo stat_test;
    stat_test.name = fbl::StringPrintf("%s/Stat", testcase.name.c_str());
    stat_test.test_fn = fbl::BindMember(&pw_op, &PathWalkOp::Stat);
    testcase.tests.push_back(std::move(stat_test));

    TestInfo unlink_test;
    unlink_test.name = fbl::StringPrintf("%s/Unlink", testcase.name.c_str());
    unlink_test.test_fn = fbl::BindMember(&pw_op, &PathWalkOp::Unlink);

    testcase.tests.push_back(std::move(unlink_test));
    testcases.push_back(std::move(testcase));
  }

  return fs_test_utils::RunTestCases(f_opts, p_opts, testcases);
}
}  // namespace fs_bench

int main(int argc, char** argv) {
  return fs_test_utils::RunWithMemFs(
      [argc, argv]() { return fs_bench::RunBenchmark(argc, argv) ? 0 : -1; });
}
